/**
 * @file DCTFFT.cuh
 * @brief Define some DCT/DST/DCST/DSCT cufft implementation functions, this file should not be included by interface header files.
 * @author zpfeng
 * @copyright (c) Copyright 2013~2050 - PANGO MICROSYSTEMS, INC.
 * ALL RIGHTS RESERVED.
 */

#ifndef __DCTFFT_CUH__
#define __DCTFFT_CUH__

#include "Utils.cuh"
#include <cufft.h>

namespace fabric_space
{
namespace dct_op
{
#define TPB (16)


inline __device__ cufftDoubleComplex ComplexMul(const cufftDoubleComplex &x, const cufftDoubleComplex &y)
{
    cufftDoubleComplex res;
    res.x = x.x * y.x - x.y * y.y;
    res.y = x.x * y.y + x.y * y.x;
    return res;
}

inline __device__ cufftComplex ComplexMul(const cufftComplex &x, const cufftComplex &y)
{
    cufftComplex res;
    res.x = x.x * y.x - x.y * y.y;
    res.y = x.x * y.y + x.y * y.x;
    return res;
}

inline __device__ cufftDoubleReal RealPartOfMul(const cufftDoubleComplex &x, const cufftDoubleComplex &y)
{
    return x.x * y.x - x.y * y.y;
}

inline __device__ cufftReal RealPartOfMul(const cufftComplex &x, const cufftComplex &y)
{
    return x.x * y.x - x.y * y.y;
}

inline __device__ cufftDoubleReal ImaginaryPartOfMul(const cufftDoubleComplex &x, const cufftDoubleComplex &y)
{
    return x.x * y.y + x.y * y.x;
}

inline __device__ cufftReal ImaginaryPartOfMul(const cufftComplex &x, const cufftComplex &y)
{
    return x.x * y.y + x.y * y.x;
}

inline __device__ cufftDoubleComplex ComplexAdd(const cufftDoubleComplex &x, const cufftDoubleComplex &y)
{
    cufftDoubleComplex res;
    res.x = x.x + y.x;
    res.y = x.y + y.y;
    return res;
}

inline __device__ cufftComplex ComplexAdd(const cufftComplex &x, const cufftComplex &y)
{
    cufftComplex res;
    res.x = x.x + y.x;
    res.y = x.y + y.y;
    return res;
}

inline __device__ cufftDoubleComplex ComplexSubtract(const cufftDoubleComplex &x, const cufftDoubleComplex &y)
{
    cufftDoubleComplex res;
    res.x = x.x - y.x;
    res.y = x.y - y.y;
    return res;
}

inline __device__ cufftComplex ComplexSubtract(const cufftComplex &x, const cufftComplex &y)
{
    cufftComplex res;
    res.x = x.x - y.x;
    res.y = x.y - y.y;
    return res;
}

inline __device__ cufftDoubleComplex ComplexConj(const cufftDoubleComplex &x)
{
    cufftDoubleComplex res;
    res.x = x.x;
    res.y = -x.y;
    return res;
}

inline __device__ cufftComplex ComplexConj(const cufftComplex &x)
{
    cufftComplex res;
    res.x = x.x;
    res.y = -x.y;
    return res;
}

inline __device__ cufftDoubleComplex ComplexMulConj(const cufftDoubleComplex &x, const cufftDoubleComplex &y)
{
    cufftDoubleComplex res;
    res.x = x.x * y.x - x.y * y.y;
    res.y = -(x.x * y.y + x.y * y.x);
    return res;
}

inline __device__ cufftComplex ComplexMulConj(const cufftComplex &x, const cufftComplex &y)
{
    cufftComplex res;
    res.x = x.x * y.x - x.y * y.y;
    res.y = -(x.x * y.y + x.y * y.x);
    return res;
}

template <typename T> void MakeCufftPlan(const int M, const int N, cufftHandle *plan)
{
}

template <> void MakeCufftPlan<cufftComplex>(const int M, const int N, cufftHandle *plan)
{
    cufftPlan2d(plan, M, N, CUFFT_R2C);
}

template <> void MakeCufftPlan<cufftDoubleComplex>(const int M, const int N, cufftHandle *plan)
{
    cufftPlan2d(plan, M, N, CUFFT_D2Z);
}

__global__ void PrecomputeExpkDCT2D(cufftComplex *expkM, cufftComplex *expkN, const int M, const int N)
{
    const int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid <= M / 2)
    {
        int hid = tid;
        cufftComplex W_h_4M = make_float2(__cosf((float)PI * hid / (2 * M)), -__sinf((float)PI * hid / (M * 2)));
        expkM[hid] = W_h_4M;
    }
    if (tid <= N / 2)
    {
        int wid = tid;
        cufftComplex W_w_4N = make_float2(__cosf((float)PI * wid / (2 * N)), -__sinf((float)PI * wid / (N * 2)));
        expkN[wid] = W_w_4N;
    }
}

__global__ void PrecomputeExpkIDCT2D(cufftComplex *expkM, cufftComplex *expkN, const int M, const int N)
{
    const int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < M)
    {
        int hid = tid;
        cufftComplex W_h_4M = make_float2(__cosf((float)PI * hid / (2 * M)), -__sinf((float)PI * hid / (M * 2)));
        expkM[hid] = W_h_4M;
    }
    if (tid <= N / 2)
    {
        int wid = tid;
        cufftComplex W_w_4N = make_float2(__cosf((float)PI * wid / (2 * N)), -__sinf((float)PI * wid / (N * 2)));
        expkN[wid] = W_w_4N;
    }
}

__global__ void PrecomputeExpkIDCT2D_v2(cufftComplex *expkM, cufftComplex *expkN, cufftComplex *expkMN_1, cufftComplex *expkMN_2, const int M, const int N)
{
    const int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < M)
    {
        int hid = tid;
        cufftComplex W_h_4M = make_float2(__cosf((float)PI * hid / (2 * M)), -__sinf((float)PI * hid / (M * 2)));
        expkM[hid] = W_h_4M;
        // expkMN_1
        cufftComplex W_h_4M_offset = make_float2(__cosf((float)PI * (hid + M) / (2 * M)), -__sinf((float)PI * (hid + M) / (M * 2)));
        expkMN_1[hid] = W_h_4M;
        expkMN_1[hid + M] = W_h_4M_offset;

        // expkMN_2
        W_h_4M = make_float2(-__sinf((float)PI * (hid - (N - 1)) / (M * 2)), -__cosf((float)PI * (hid - (N - 1)) / (2 * M)));
        W_h_4M_offset = make_float2(-__sinf((float)PI * (hid - (N - 1) + M) / (M * 2)), -__cosf((float)PI * (hid - (N - 1) + M) / (2 * M)));
        expkMN_2[hid] = W_h_4M;
        expkMN_2[hid + M] = W_h_4M_offset;
    }
    if (tid <= N / 2)
    {
        int wid = tid;
        cufftComplex W_w_4N = make_float2(__cosf((float)PI * wid / (2 * N)), -__sinf((float)PI * wid / (N * 2)));
        expkN[wid] = W_w_4N;
    }
}

__global__ void PrecomputeExpkDSCT2D(cufftComplex *expkM, cufftComplex *expkN, const int M, const int N)
{
    const int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < M)
    {
        int hid = tid;
        cufftComplex W_h_4M = make_float2(__cosf((float)PI * hid / (2 * M)), -__sinf((float)PI * hid / (M * 2)));
        expkM[hid] = W_h_4M;
    }
    if (tid <= N / 2)
    {
        int wid = tid;
        cufftComplex W_w_4N = make_float2(__cosf((float)PI * wid / (2 * N)), -__sinf((float)PI * wid / (N * 2)));
        expkN[wid] = W_w_4N;
    }
}

__global__ void PrecomputeExpkDCST2D(cufftComplex *expkM, cufftComplex *expkN, const int M, const int N)
{
    const int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < M)
    {
        int hid = tid;
        cufftComplex W_h_4M = make_float2(__cosf((float)PI * hid / (2 * M)), -__sinf((float)PI * hid / (M * 2)));
        expkM[hid] = W_h_4M;
    }
    if (tid <= N / 2)
    {
        int wid = tid;
        cufftComplex W_w_4N = make_float2(__cosf((float)PI * wid / (2 * N)), -__sinf((float)PI * wid / (N * 2)));
        expkN[wid] = W_w_4N;
    }
}

template <typename T> __global__ void dct2d_preprocess(const T *x, T *y, const int M, const int N, const int halfN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < M && wid < N)
    {
        int index;
        int cond = (((hid & 1) == 0) << 1) | ((wid & 1) == 0);
        switch (cond)
        {
            case 0:
            {
                index = INDEX((M << 1) - (hid + 1), N - ((wid + 1) >> 1), halfN);
                break;
            }
            case 1:
            {
                index = INDEX((M << 1) - (hid + 1), (wid >> 1), halfN);
                break;
            }
            case 2:
            {
                index = INDEX(hid, N - ((wid + 1) >> 1), halfN);
                break;
            }
            case 3:
            {
                index = INDEX(hid, (wid >> 1), halfN);
                break;
            }
            default:
            {
                break;
            }
        }
        y[index] = x[INDEX(hid, wid, N)];
    }
}

template <typename T, typename TComplex>
__global__ void __launch_bounds__(TPB *TPB, 8) dct2d_postprocess(const TComplex *V, T *y, const int M, const int N, const int halfM, const int halfN, const T two_over_MN,
    const T four_over_MN, const TComplex *__restrict__ expkM, const TComplex *__restrict__ expkN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < halfM && wid < halfN)
    {
        int cond = ((hid != 0) << 1) | (wid != 0);
        switch (cond)
        {
            case 0:
            {
                y[0] = V[0].x * four_over_MN;
                y[halfN] = RealPartOfMul(expkN[halfN], V[halfN]) * four_over_MN;
                y[INDEX(halfM, 0, N)] = expkM[halfM].x * V[INDEX(halfM, 0, halfN + 1)].x * four_over_MN;
                y[INDEX(halfM, halfN, N)] = expkM[halfM].x * RealPartOfMul(expkN[halfN], V[INDEX(halfM, halfN, halfN + 1)]) * four_over_MN;
                break;
            }
            case 1:
            {
                TComplex tmp;
                tmp = V[wid];
                y[wid] = RealPartOfMul(expkN[wid], tmp) * four_over_MN;
                y[N - wid] = -ImaginaryPartOfMul(expkN[wid], tmp) * four_over_MN;
                tmp = V[INDEX(halfM, wid, halfN + 1)];
                y[INDEX(halfM, wid, N)] = expkM[halfM].x * RealPartOfMul(expkN[wid], tmp) * four_over_MN;
                y[INDEX(halfM, N - wid, N)] = -expkM[halfM].x * ImaginaryPartOfMul(expkN[wid], tmp) * four_over_MN;
                break;
            }
            case 2:
            {
                TComplex tmp1, tmp2, tmp_up, tmp_down;
                tmp1 = V[INDEX(hid, 0, halfN + 1)];
                tmp2 = V[INDEX(M - hid, 0, halfN + 1)];
                tmp_up.x = expkM[hid].x * (tmp1.x + tmp2.x) + expkM[hid].y * (tmp2.y - tmp1.y);
                tmp_down.x = -expkM[hid].y * (tmp1.x + tmp2.x) + expkM[hid].x * (tmp2.y - tmp1.y);
                y[INDEX(hid, 0, N)] = tmp_up.x * two_over_MN;
                y[INDEX(M - hid, 0, N)] = tmp_down.x * two_over_MN;
                tmp1 = ComplexAdd(V[INDEX(hid, halfN, halfN + 1)], V[INDEX(M - hid, halfN, halfN + 1)]);
                tmp2 = ComplexSubtract(V[INDEX(hid, halfN, halfN + 1)], V[INDEX(M - hid, halfN, halfN + 1)]);
                tmp_up.x = expkM[hid].x * tmp1.x - expkM[hid].y * tmp2.y;
                tmp_up.y = expkM[hid].x * tmp1.y + expkM[hid].y * tmp2.x;
                tmp_down.x = -expkM[hid].y * tmp1.x - expkM[hid].x * tmp2.y;
                tmp_down.y = -expkM[hid].y * tmp1.y + expkM[hid].x * tmp2.x;
                y[INDEX(hid, halfN, N)] = RealPartOfMul(expkN[halfN], tmp_up) * two_over_MN;
                y[INDEX(M - hid, halfN, N)] = RealPartOfMul(expkN[halfN], tmp_down) * two_over_MN;
                break;
            }
            case 3:
            {
                TComplex tmp1, tmp2, tmp_up, tmp_down;
                tmp1 = ComplexAdd(V[INDEX(hid, wid, halfN + 1)], V[INDEX(M - hid, wid, halfN + 1)]);
                tmp2 = ComplexSubtract(V[INDEX(hid, wid, halfN + 1)], V[INDEX(M - hid, wid, halfN + 1)]);
                tmp_up.x = expkM[hid].x * tmp1.x - expkM[hid].y * tmp2.y;
                tmp_up.y = expkM[hid].x * tmp1.y + expkM[hid].y * tmp2.x;
                tmp_down.x = -expkM[hid].y * tmp1.x - expkM[hid].x * tmp2.y;
                tmp_down.y = -expkM[hid].y * tmp1.y + expkM[hid].x * tmp2.x;
                y[INDEX(hid, wid, N)] = RealPartOfMul(expkN[wid], tmp_up) * two_over_MN;
                y[INDEX(M - hid, wid, N)] = RealPartOfMul(expkN[wid], tmp_down) * two_over_MN;
                y[INDEX(hid, N - wid, N)] = -ImaginaryPartOfMul(expkN[wid], tmp_up) * two_over_MN;
                y[INDEX(M - hid, N - wid, N)] = -ImaginaryPartOfMul(expkN[wid], tmp_down) * two_over_MN;
                break;
            }
            default:
            {
                assert(0);
                break;
            }
        }
    }
}

template <typename T, typename TComplex>
__global__ void __launch_bounds__(TPB *TPB, 8) idct2d_preprocess(
    const T *input, TComplex *output, const int M, const int N, const int halfM, const int halfN, const TComplex *__restrict__ expkM, const TComplex *__restrict__ expkN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < halfM && wid < halfN)
    {
        int cond = ((hid != 0) << 1) | (wid != 0);
        switch (cond)
        {
            case 0:
            {
                T tmp1;
                TComplex tmp_up;
                output[0].x = input[0];
                output[0].y = 0;
                tmp1 = input[halfN];
                tmp_up.x = tmp1;
                tmp_up.y = tmp1;
                output[halfN] = ComplexConj(ComplexMul(expkN[halfN], tmp_up));
                tmp1 = input[INDEX(halfM, 0, N)];
                tmp_up.x = tmp1;
                tmp_up.y = tmp1;
                output[INDEX(halfM, 0, halfN + 1)] = ComplexConj(ComplexMul(expkM[halfM], tmp_up));
                tmp1 = input[INDEX(halfM, halfN, N)];
                tmp_up.x = 0;
                tmp_up.y = 2 * tmp1;
                output[INDEX(halfM, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[halfM], expkN[halfN]), tmp_up));
                break;
            }
            case 1:
            {
                TComplex tmp_up;
                tmp_up.x = input[wid];
                tmp_up.y = input[N - wid];
                output[wid] = ComplexConj(ComplexMul(expkN[wid], tmp_up));
                T tmp1 = input[INDEX(halfM, wid, N)];
                T tmp2 = input[INDEX(halfM, N - wid, N)];
                tmp_up.x = tmp1 - tmp2;
                tmp_up.y = tmp1 + tmp2;
                output[INDEX(halfM, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[halfM], expkN[wid]), tmp_up));
                break;
            }
            case 2:
            {
                T tmp1, tmp3;
                TComplex tmp_up, tmp_down;
                tmp1 = input[INDEX(hid, 0, N)];
                tmp3 = input[INDEX(M - hid, 0, N)];
                tmp_up.x = tmp1;
                tmp_up.y = tmp3;
                tmp_down.x = tmp3;
                tmp_down.y = tmp1;
                output[INDEX(hid, 0, halfN + 1)] = ComplexConj(ComplexMul(expkM[hid], tmp_up));
                output[INDEX(M - hid, 0, halfN + 1)] = ComplexConj(ComplexMul(expkM[M - hid], tmp_down));
                tmp1 = input[INDEX(hid, halfN, N)];
                tmp3 = input[INDEX(M - hid, halfN, N)];
                tmp_up.x = tmp1 - tmp3;
                tmp_up.y = tmp3 + tmp1;
                tmp_down.x = tmp3 - tmp1;
                tmp_down.y = tmp1 + tmp3;
                output[INDEX(hid, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[hid], expkN[halfN]), tmp_up));
                output[INDEX(M - hid, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[M - hid], expkN[halfN]), tmp_down));
                break;
            }
            case 3:
            {
                T tmp1 = input[INDEX(hid, wid, N)];
                T tmp2 = input[INDEX(hid, N - wid, N)];
                T tmp3 = input[INDEX(M - hid, wid, N)];
                T tmp4 = input[INDEX(M - hid, N - wid, N)];
                TComplex tmp_up, tmp_down;
                tmp_up.x = tmp1 - tmp4;
                tmp_up.y = tmp3 + tmp2;
                tmp_down.x = tmp3 - tmp2;
                tmp_down.y = tmp1 + tmp4;

                output[INDEX(hid, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[hid], expkN[wid]), tmp_up));
                output[INDEX(M - hid, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[M - hid], expkN[wid]), tmp_down));
                break;
            }
            default:
                assert(0);
                break;
        }
    }
}

template <typename T, typename TComplex>
__global__ void __launch_bounds__(TPB *TPB, 8) idct2d_preprocess_v2(const T *input, TComplex *output, const int M, const int N, const int halfM, const int halfN,
    const TComplex *__restrict__ expkM, const TComplex *__restrict__ expkN, const TComplex *__restrict__ expkMN_1, const TComplex *__restrict__ expkMN_2)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < halfM && wid < halfN)
    {
        int cond = ((hid != 0) << 1) | (wid != 0);
        switch (cond)
        {
            case 0:
            {
                T tmp1;
                TComplex tmp_up;
                output[0].x = input[0];
                output[0].y = 0;
                tmp1 = input[halfN];
                tmp_up.x = tmp1;
                tmp_up.y = tmp1;
                output[halfN] = ComplexMulConj(expkN[halfN], tmp_up);
                tmp1 = input[INDEX(halfM, 0, N)];
                tmp_up.x = tmp1;
                tmp_up.y = tmp1;
                output[INDEX(halfM, 0, halfN + 1)] = ComplexMulConj(expkM[halfM], tmp_up);
                tmp1 = input[INDEX(halfM, halfN, N)];
                tmp_up.x = 0;
                tmp_up.y = 2 * tmp1;
                output[INDEX(halfM, halfN, halfN + 1)] = ComplexMulConj(expkMN_1[halfM + halfN], tmp_up);
                break;
            }
            case 1:
            {
                TComplex tmp_up;
                tmp_up.x = input[wid];
                tmp_up.y = input[N - wid];
                output[wid] = ComplexMulConj(expkN[wid], tmp_up);
                T tmp1 = input[INDEX(halfM, wid, N)];
                T tmp2 = input[INDEX(halfM, N - wid, N)];
                tmp_up.x = tmp1 - tmp2;
                tmp_up.y = tmp1 + tmp2;
                output[INDEX(halfM, wid, halfN + 1)] = ComplexMulConj(expkMN_1[halfM + wid], tmp_up);
                break;
            }
            case 2:
            {
                T tmp1, tmp3;
                TComplex tmp_up, tmp_down;
                tmp1 = input[INDEX(hid, 0, N)];
                tmp3 = input[INDEX(M - hid, 0, N)];
                tmp_down.x = tmp3;
                tmp_down.y = tmp1;
                // two outputs are conjugate
                tmp_up = ComplexMul(expkM[M - hid], tmp_down);
                output[INDEX(hid, 0, halfN + 1)] = tmp_up;
                output[INDEX(M - hid, 0, halfN + 1)] = ComplexConj(tmp_up);
                tmp1 = input[INDEX(hid, halfN, N)];
                tmp3 = input[INDEX(M - hid, halfN, N)];
                tmp_up.x = tmp1 - tmp3;
                tmp_up.y = tmp3 + tmp1;
                tmp_down.x = tmp3 - tmp1;
                tmp_down.y = tmp1 + tmp3;
                output[INDEX(hid, halfN, halfN + 1)] = ComplexMulConj(expkMN_1[hid + halfN], tmp_up);
                output[INDEX(M - hid, halfN, halfN + 1)] = ComplexMulConj(expkMN_2[halfN - hid + (N - 1)], tmp_down);
                break;
            }
            case 3:
            {
                T tmp1 = input[INDEX(hid, wid, N)];
                T tmp2 = input[INDEX(hid, N - wid, N)];
                T tmp3 = input[INDEX(M - hid, wid, N)];
                T tmp4 = input[INDEX(M - hid, N - wid, N)];
                TComplex tmp_up, tmp_down;
                tmp_up.x = tmp1 - tmp4;
                tmp_up.y = tmp3 + tmp2;
                tmp_down.x = tmp3 - tmp2;
                tmp_down.y = tmp1 + tmp4;
                output[INDEX(hid, wid, halfN + 1)] = ComplexMulConj(expkMN_1[hid + wid], tmp_up);
                output[INDEX(M - hid, wid, halfN + 1)] = ComplexMulConj(expkMN_2[wid - hid + (N - 1)], tmp_down);
                break;
            }
            default:
                assert(0);
                break;
        }
    }
}

template <typename T> __global__ void idct2d_postprocess(const T *x, T *y, const int M, const int N, const int halfN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < M && wid < N)
    {
        int cond = ((hid < M / 2) << 1) | (wid < N / 2);
        int index;
        switch (cond)
        {
            case 0:
                index = INDEX(((M - hid) << 1) - 1, ((N - wid) << 1) - 1, N);
                break;
            case 1:
                index = INDEX(((M - hid) << 1) - 1, wid << 1, N);
                break;
            case 2:
                index = INDEX(hid << 1, ((N - wid) << 1) - 1, N);
                break;
            case 3:
                index = INDEX(hid << 1, wid << 1, N);
                break;
            default:
                assert(0);
                break;
        }
        y[index] = x[INDEX(hid, wid, N)];
    }
}

// Adpated from idct2d_preprocess(). The only change is the reordered input
// if (hid != 0)
//     new_input[hid][wid] = input[M - hid][wid];
// else
//     new_input[0][wid] = 0
template <typename T, typename TComplex>
__global__ void __launch_bounds__(TPB *TPB, 8) idxst_idct_preprocess(
    const T *input, TComplex *output, const int M, const int N, const int halfM, const int halfN, const TComplex *__restrict__ expkM, const TComplex *__restrict__ expkN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < halfM && wid < halfN)
    {
        int cond = ((hid != 0) << 1) | (wid != 0);
        switch (cond)
        {
            case 0:
            {
                T tmp1;
                TComplex tmp_up;
                output[0].x = 0;
                output[0].y = 0;
                output[halfN].x = 0;
                output[halfN].y = 0;
                tmp1 = input[INDEX(halfM, 0, N)];
                tmp_up.x = tmp1;
                tmp_up.y = tmp1;
                output[INDEX(halfM, 0, halfN + 1)] = ComplexConj(ComplexMul(expkM[halfM], tmp_up));
                tmp1 = input[INDEX(halfM, halfN, N)];
                tmp_up.x = 0;
                tmp_up.y = 2 * tmp1;
                output[INDEX(halfM, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[halfM], expkN[halfN]), tmp_up));
                break;
            }
            case 1:
            {
                output[wid].x = 0;
                output[wid].y = 0;
                TComplex tmp_up;
                T tmp1 = input[INDEX(halfM, wid, N)];
                T tmp2 = input[INDEX(halfM, N - wid, N)];
                tmp_up.x = tmp1 - tmp2;
                tmp_up.y = tmp1 + tmp2;
                output[INDEX(halfM, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[halfM], expkN[wid]), tmp_up));
                break;
            }
            case 2:
            {
                T tmp1, tmp3;
                TComplex tmp_up, tmp_down;
                tmp1 = input[INDEX(M - hid, 0, N)];
                tmp3 = input[INDEX(hid, 0, N)];
                tmp_up.x = tmp1;
                tmp_up.y = tmp3;
                tmp_down.x = tmp3;
                tmp_down.y = tmp1;
                output[INDEX(hid, 0, halfN + 1)] = ComplexConj(ComplexMul(expkM[hid], tmp_up));
                output[INDEX(M - hid, 0, halfN + 1)] = ComplexConj(ComplexMul(expkM[M - hid], tmp_down));
                tmp1 = input[INDEX(M - hid, halfN, N)];
                tmp3 = input[INDEX(hid, halfN, N)];
                tmp_up.x = tmp1 - tmp3;
                tmp_up.y = tmp3 + tmp1;
                tmp_down.x = tmp3 - tmp1;
                tmp_down.y = tmp1 + tmp3;
                output[INDEX(hid, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[hid], expkN[halfN]), tmp_up));
                output[INDEX(M - hid, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[M - hid], expkN[halfN]), tmp_down));
                break;
            }
            case 3:
            {
                T tmp1 = input[INDEX(M - hid, wid, N)];
                T tmp2 = input[INDEX(M - hid, N - wid, N)];
                T tmp3 = input[INDEX(hid, wid, N)];
                T tmp4 = input[INDEX(hid, N - wid, N)];
                TComplex tmp_up, tmp_down;
                tmp_up.x = tmp1 - tmp4;
                tmp_up.y = tmp3 + tmp2;
                tmp_down.x = tmp3 - tmp2;
                tmp_down.y = tmp1 + tmp4;
                output[INDEX(hid, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[hid], expkN[wid]), tmp_up));
                output[INDEX(M - hid, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[M - hid], expkN[wid]), tmp_down));
                break;
            }
            default:
            {
                assert(0);
                break;
            }
        }
    }
}

// Adpated from idct2d_postprocess() with changes on sign and scale
// if (hid % 2 == 1)
//     new_output[hid][wid] = -output[hid][wid];
// else
//     new_output[hid][wid] = output[hid][wid];
template <typename T> __global__ void idxst_idct_postprocess(const T *x, T *y, const int M, const int N, const int halfN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < M && wid < N)
    {
        int cond = ((hid < M / 2) << 1) | (wid < N / 2);
        int index;
        switch (cond)
        {
            case 0:
            {
                index = INDEX(((M - hid) << 1) - 1, ((N - wid) << 1) - 1, N);
                y[index] = -x[INDEX(hid, wid, N)];
                break;
            }
            case 1:
            {
                index = INDEX(((M - hid) << 1) - 1, wid << 1, N);
                y[index] = -x[INDEX(hid, wid, N)];
                break;
            }
            case 2:
            {
                index = INDEX(hid << 1, ((N - wid) << 1) - 1, N);
                y[index] = x[INDEX(hid, wid, N)];
                break;
            }
            case 3:
            {
                index = INDEX(hid << 1, wid << 1, N);
                y[index] = x[INDEX(hid, wid, N)];
                break;
            }
            default:
            {
                assert(0);
                break;
            }
        }
    }
}

// Adpated from idct2d_preprocess(). The only change is the reordered input
// if (wid != 0)
//     new_input[hid][wid] = input[hid][N - wid];
// else
//     new_input[hid][0] = 0
template <typename T, typename TComplex>
__global__ void __launch_bounds__(TPB *TPB, 8) idct_idxst_preprocess(
    const T *input, TComplex *output, const int M, const int N, const int halfM, const int halfN, const TComplex *__restrict__ expkM, const TComplex *__restrict__ expkN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < halfM && wid < halfN)
    {
        int cond = ((hid != 0) << 1) | (wid != 0);
        switch (cond)
        {
            case 0:
            {
                T tmp1;
                TComplex tmp_up;
                output[0].x = 0;
                output[0].y = 0;
                tmp1 = input[halfN];
                tmp_up.x = tmp1;
                tmp_up.y = tmp1;
                output[halfN] = ComplexConj(ComplexMul(expkN[halfN], tmp_up));
                output[INDEX(halfM, 0, halfN + 1)].x = 0;
                output[INDEX(halfM, 0, halfN + 1)].y = 0;
                tmp1 = input[INDEX(halfM, halfN, N)];
                tmp_up.x = 0;
                tmp_up.y = 2 * tmp1;
                output[INDEX(halfM, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[halfM], expkN[halfN]), tmp_up));
                break;
            }
            case 1:
            {
                TComplex tmp_up;
                tmp_up.x = input[N - wid];
                tmp_up.y = input[wid];
                output[wid] = ComplexConj(ComplexMul(expkN[wid], tmp_up));
                T tmp1 = input[INDEX(halfM, N - wid, N)];
                T tmp2 = input[INDEX(halfM, wid, N)];
                tmp_up.x = tmp1 - tmp2;
                tmp_up.y = tmp1 + tmp2;
                output[INDEX(halfM, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[halfM], expkN[wid]), tmp_up));
                break;
            }
            case 2:
            {
                T tmp1, tmp3;
                TComplex tmp_up, tmp_down;
                output[INDEX(hid, 0, halfN + 1)].x = 0;
                output[INDEX(hid, 0, halfN + 1)].y = 0;
                output[INDEX(M - hid, 0, halfN + 1)].x = 0;
                output[INDEX(M - hid, 0, halfN + 1)].y = 0;
                tmp1 = input[INDEX(hid, halfN, N)];
                tmp3 = input[INDEX(M - hid, halfN, N)];
                tmp_up.x = tmp1 - tmp3;
                tmp_up.y = tmp3 + tmp1;
                tmp_down.x = tmp3 - tmp1;
                tmp_down.y = tmp1 + tmp3;
                output[INDEX(hid, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[hid], expkN[halfN]), tmp_up));
                output[INDEX(M - hid, halfN, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[M - hid], expkN[halfN]), tmp_down));
                break;
            }
            case 3:
            {
                T tmp1 = input[INDEX(hid, N - wid, N)];
                T tmp2 = input[INDEX(hid, wid, N)];
                T tmp3 = input[INDEX(M - hid, N - wid, N)];
                T tmp4 = input[INDEX(M - hid, wid, N)];
                TComplex tmp_up, tmp_down;
                tmp_up.x = tmp1 - tmp4;
                tmp_up.y = tmp3 + tmp2;
                tmp_down.x = tmp3 - tmp2;
                tmp_down.y = tmp1 + tmp4;
                output[INDEX(hid, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[hid], expkN[wid]), tmp_up));
                output[INDEX(M - hid, wid, halfN + 1)] = ComplexConj(ComplexMul(ComplexMul(expkM[M - hid], expkN[wid]), tmp_down));
                break;
            }
            default:
            {
                assert(0);
                break;
            }
        }
    }
}

// Adpated from idct2d_postprocess() with changes on sign and scale
// if (wid % 2 == 1)
//     new_output[hid][wid] = -output[hid][wid];
// else
//     new_output[hid][wid] = output[hid][wid];
template <typename T> __global__ void idct_idxst_postprocess(const T *x, T *y, const int M, const int N, const int halfN)
{
    const int wid = blockDim.x * blockIdx.x + threadIdx.x;
    const int hid = blockDim.y * blockIdx.y + threadIdx.y;
    if (hid < M && wid < N)
    {
        int cond = ((hid < M / 2) << 1) | (wid < N / 2);
        int index;
        switch (cond)
        {
            case 0:
            {
                index = INDEX(((M - hid) << 1) - 1, ((N - wid) << 1) - 1, N);
                y[index] = -x[INDEX(hid, wid, N)];
                break;
            }
            case 1:
            {
                index = INDEX(((M - hid) << 1) - 1, wid << 1, N);
                y[index] = x[INDEX(hid, wid, N)];
                break;
            }
            case 2:
            {
                index = INDEX(hid << 1, ((N - wid) << 1) - 1, N);
                y[index] = -x[INDEX(hid, wid, N)];
                break;
            }
            case 3:
            {
                index = INDEX(hid << 1, wid << 1, N);
                y[index] = x[INDEX(hid, wid, N)];
                break;
            }
            default:
            {
                assert(0);
                break;
            }
        }
    }
}

void fft2D(cufftReal *d_x, cufftComplex *d_y, cufftHandle &plan)
{
    cufftExecR2C(plan, d_x, d_y);
    cudaDeviceSynchronize();
}

void ifft2D(cufftComplex *d_x, cufftReal *d_y, cufftHandle &plan)
{
    cufftExecC2R(plan, d_x, d_y);
    cudaDeviceSynchronize();
}

}  //!< end of namespace dct_op
}  //!< end of namespace fabric_space

#endif
